from matplotlib import pyplot as plt
from ColorAndStyle import COLOR, STYLE


class DrawScatter:
    def __init__(self):
        self._title_size = 30
        self._axis_size = 30
        self._figure = plt.figure(2)
        self._ax = self._figure.gca()
        self._ax.set_xlim(-100, 1100)
        self._ax.set_ylim(40, 120)

    def _draw_scatter(self, x, y, color=COLOR.MAGENTA, marker=STYLE.star,
                      label=None, title=None, infor='the last improved step', special_index=None):
        plt.axis([min(x) - 1, max(x) + 1, min(y) - 0.1, max(y) + 0.1])
        plt.scatter(x, y, color=color, marker=marker, label=label)
        plt.title(title, fontsize=23)
        if special_index:
            plt.annotate(infor, xy=(x[special_index - 1], y[special_index - 1]),
                         xytext=(x[special_index - 1], y[special_index - 1] + 0.2),
                         arrowprops=dict(facecolor='black', shrink=0.05), size=self._title_size)
        plt.tick_params(labelsize=self._axis_size)

    def draw_scatter(self, x, y, color=COLOR.MAGENTA, marker=STYLE.star,
                     label=None, title=None, infor='the last improved step', special_index=None):
        self._draw_scatter(x, y, color, marker,
                           label, title, infor, special_index)
        plt.xlabel('Iteration', fontsize=self._title_size, horizontalalignment='center')
        plt.ylabel('Objective', fontsize=self._title_size)
        plt.legend()
        plt.show()

    def draw_scatters(self, xs, ys, labels=None, title=None, infor='the last improved step', special_index=None):
        for x, y, color, label in zip(xs, ys, COLOR.COLORS[6:], labels):
            self._draw_scatter(x, y, color, STYLE.star, label, title, infor, special_index)
        plt.xlabel('Iteration', fontsize=self._title_size, horizontalalignment='center')
        plt.ylabel('Objective', fontsize=self._title_size)
        plt.legend(fontsize=self._title_size)
        plt.show()


if __name__ == "__main__":
    s = DrawScatter()
    s.draw_scatter([1, 2, 3], [4, 5, 6], special_index=2)
